# Copyright (c) 2022 Presto Labs Pte. Ltd.
# Author: zijin

import json
import warnings
import heapq

from absl import app, flags
from python.experimental.jingyuan.exam_model import norm_fi, process_feature_name

FLAGS = flags.FLAGS
DEFAULT_DS = 30


def main(argv):
  assert FLAGS.input_file
  df = norm_fi(regex=FLAGS.regex, path=FLAGS.path)
  fidict = {name: score for name, score in zip(df.index, df['fis'].values)}
  fjson = json.load(open(FLAGS.input_file, 'r'))
  print('=== before: ', len(fjson['features']))
  feature_dict = {}
  # out of dict features
  ood_features = []
  for name in fjson['features']:
    keyname = process_feature_name(name)
    if keyname in fidict:
      feature_dict[name] = fidict[keyname]
    else:
      ood_features.append(name)
      warnings.warn(f"Feature {name} not in dict")

  if FLAGS.zscore != None and FLAGS.features_count == None:
    new_features = [fn for fn in feature_dict.keys() if feature_dict[fn] > FLAGS.zscore]
  elif FLAGS.zscore == None and FLAGS.features_count != None:
    # get top X zscore features
    new_features = heapq.nlargest(FLAGS.features_count, feature_dict, key=feature_dict.get)
  else:
    warnings.warn("Please specify zscore or feature count threshold.")
    return

  removed_features = [fn for fn in feature_dict.keys() if fn not in new_features]
  print('=== features removed: ', removed_features)
  print('=== features not in zscore dict and just keep: ', len(ood_features), ood_features)
  print('=== after: ', len(new_features + ood_features))
  new_fjson = {'features': new_features + ood_features}
  if FLAGS.output_file:
    json.dump(new_fjson, open(FLAGS.output_file, 'w'), indent=2)
  else:
    json.dump(new_fjson, open(FLAGS.input_file, 'w'), indent=2)


# ./pyrunner python/experimental/zijin/filter_feature.py --zscore=-2 --regex=model_smm_binance_perp
if __name__ == '__main__':
  flags.DEFINE_string(
      'input_file',
      '/remote/iosg/home-3/zijin/workspace/coin/tmp/model_0930_bns/features.json',
      '')
  flags.DEFINE_string(
      'output_file',
      '/remote/iosg/home-3/zijin/workspace/coin/tmp/model_0930_bns/filtered_features.json',
      '')
  flags.DEFINE_float('zscore', None, '')
  flags.DEFINE_string('path', None, '')
  flags.DEFINE_string('regex', None, '')
  flags.DEFINE_integer('features_count', None, 'Keep X features sorted by zscore')
  app.run(main)
